"""
Adapted from salesforce@LAVIS. Below is the original copyright:
 Copyright (c) 2023, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE_Lavis file in the repo root or https://opensource.org/licenses/BSD-3-Clause
"""
import contextlib
import logging
import os
import time
import datetime

import torch
import torch.nn as nn
import torch.distributed as dist
import torch.nn.functional as F

import timechat.common.dist_utils as dist_utils
from timechat.common.dist_utils import download_cached_file
from timechat.common.utils import is_url
from timechat.common.logger import MetricLogger
from timechat.models.base_model import BaseModel
from timechat.models.Qformer import BertConfig, BertLMHeadModel
from timechat.models.eva_vit import create_eva_vit_g
from transformers import BertTokenizer


class Blip2Base(BaseModel):
    @classmethod
    def init_tokenizer(cls): #初始化分词器
        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") #使用预训练的bert-base-uncased模型初始化BertTokenizer
        tokenizer.add_special_tokens({"bos_token": "[DEC]"}) #向分词器中添加一个特殊token“[DEC]”，作为解码的开始标志
        return tokenizer #返回初始化好的分词器对象

    # 用于管理自动混合精度计算的上下文
    def maybe_autocast(self, dtype=torch.float16):
        # if on cpu, don't use autocast
        # if on gpu, use autocast with dtype if provided, otherwise use torch.float16
        enable_autocast = self.device != torch.device("cpu") #当前设备如果不是cpu，则启用自动混合精度计算

        if enable_autocast: #如果启用
            return torch.cuda.amp.autocast(dtype=dtype) #返回一个自动混合精度上下文管理器，数据类型为dtype，用以在GPU加速训练
        else:
            return contextlib.nullcontext() #返回一个空的上下文管理器

    @classmethod
    def init_Qformer(cls, num_query_token, vision_width, cross_attention_freq=2): #初始化Qformer
        encoder_config = BertConfig.from_pretrained("bert-base-uncased") #从预训练的模型中加载BertConfig配置，作为Qf的基础配置
        encoder_config.encoder_width = vision_width #设置编码器的宽度，使其与视觉编码器的输出特征维度相匹配
        # insert cross-attention layer every other block
        encoder_config.add_cross_attention = True #启用交叉注意力机制，使Qf能关注视觉编码器的输出
        encoder_config.cross_attention_freq = cross_attention_freq #设置交叉注意力的频率，默认为2
        encoder_config.query_length = num_query_token #设置查询token的数量
        Qformer = BertLMHeadModel(config=encoder_config) #根据修改后的配置创建一个BertLMHeadModel实例，即Qf模型
        query_tokens = nn.Parameter(
            torch.zeros(1, num_query_token, encoder_config.hidden_size)
        ) #创建一个可学习的查询token参数，用于初始化查询向量
        query_tokens.data.normal_(mean=0.0, std=encoder_config.initializer_range) #对查询 token 参数进行正态分布初始化，均值为 0.0，标准差为配置中的初始化范围，有助于模型训练的稳定性和收敛性。
        return Qformer, query_tokens

    @classmethod
    #初始化视觉编码器
    def init_vision_encoder(
        cls, url_or_filename, img_size, drop_path_rate, use_grad_checkpoint, precision
    ):
        assert "eva_vit_g" in url_or_filename, "vit model must be eva_vit_g for current version of MiniGPT-4"
        visual_encoder = create_eva_vit_g(
            url_or_filename, img_size, drop_path_rate, use_grad_checkpoint, precision
        ) #创建eva_vit_g视觉编码器，传入模型的URL或文件名，图像大小，丢弃路径率，是否使用梯度检查点以及精度等参数，构建用于处理视觉信息的模型

        ln_vision = LayerNorm(visual_encoder.num_features) #创建一个LayerNorm层，对视觉编码器的输出特征进行归一化处理
        return visual_encoder, ln_vision

    #用于从预训练模型文件或URL加载模型参数
    def load_from_pretrained(self, url_or_filename):
        if is_url(url_or_filename): #判断输入是否为URL
            cached_file = download_cached_file(
                url_or_filename, check_hash=False, progress=True
            ) #如果是URL，调用函数下载文件并缓存到本地，check_hash=False, progress=True表示不检查哈希值，显示下载进度
            checkpoint = torch.load(cached_file, map_location="cpu") #从缓存文件中加载模型检查点，将其存储在cpu上
        elif os.path.isfile(url_or_filename): #若输入为文件路径且存在
            checkpoint = torch.load(url_or_filename, map_location="cpu") #直接从文件路径加载模型检查点到cpu
        else:
            raise RuntimeError("checkpoint url or path is invalid") #否则抛出错误

        state_dict = checkpoint["model"] #从检查点提前模型的状态字典，包含模型的参数和状态信息

        msg = self.load_state_dict(state_dict, strict=False) #加载状态字典到当前模型中，strict=False表示允许模型中存在于状态字典不完全匹配的参数

        # logging.info("Missing keys {}".format(msg.missing_keys))
        logging.info("load checkpoint from %s" % url_or_filename) #记录缺失的键信息，即模型中存在但状态字典中不存在的参数

        return msg

#用于禁用模型训练模型的切换
def disabled_train(self, mode=True):
    """Overwrite model.train with this function to make sure train/eval mode
    does not change anymore."""
    return self


class LayerNorm(nn.LayerNorm):
    """Subclass torch's LayerNorm to handle fp16."""

    def forward(self, x: torch.Tensor):
        orig_type = x.dtype  #保存输入张量x的原始数据类型
        ret = super().forward(x.type(torch.float32)) #先将输入张量的数据类型转换为float32
        return ret.type(orig_type)

#用于计算模型在数据加载器data_loader上的文本-图像相似性矩阵
def compute_sim_matrix(model, data_loader, **kwargs):
    k_test = kwargs.pop("k_test") #从kwargs中取出用于测试的k值

    metric_logger = MetricLogger(delimiter="  ") #创建一个MetricLogger对象，用于记录和显示指标信息，delimiter指定分隔符
    header = "Evaluation:" #设置日志的头部信息，表示这是在评估阶段的日志

    logging.info("Computing features for evaluation...") #记录一条信息日志，表明正在计算评估所需的特征
    start_time = time.time() #记录开始时间

    texts = data_loader.dataset.text #获取数据加载器中数据集的文本数据
    num_text = len(texts) #计算文本数据的数量
    text_bs = 256 #设置文本批次大小为256
    text_ids = [] #用于存储文本的输入id
    text_embeds = [] #存储文本的嵌入向量
    text_atts = [] #存储文本的注意力掩码
    for i in range(0, num_text, text_bs): #循环遍历文本数据，按批次大小进行分割
        text = texts[i : min(num_text, i + text_bs)] #获取当前批次的文本数据
        text_input = model.tokenizer(
            text,
            padding="max_length",
            truncation=True,
            max_length=35,
            return_tensors="pt",
        ).to(model.device) #对文本进行分词和编码，生成输入id，注意力掩码等张量
        text_feat = model.forward_text(text_input) #将文本输入传递给模型的文本编码部分，获取文本特征
        text_embed = F.normalize(model.text_proj(text_feat)) #对文本特征进行投影，并进行归一化处理，使其具有单位长度
        text_embeds.append(text_embed) #将当前批次的文本嵌入向量添加到列表中
        text_ids.append(text_input.input_ids) #添加输入id
        text_atts.append(text_input.attention_mask) #添加注意力掩码

    text_embeds = torch.cat(text_embeds, dim=0) #将所有批次的文本嵌入向量在第0维上拼接成一个整体的嵌入张量
    text_ids = torch.cat(text_ids, dim=0) #同上
    text_atts = torch.cat(text_atts, dim=0) #同上

    vit_feats = [] #用于存储图像的vit特征
    image_embeds = [] #存储图像的嵌入向量
    for samples in data_loader:
        image = samples["image"] #获取当前批次的图像数据

        image = image.to(model.device) #将图像数据移动到模型所在的设备上
        image_feat, vit_feat = model.forward_image(image) #将图像输入模型传递给模型的图像编码部分，获取图像特征和vit特征
        image_embed = model.vision_proj(image_feat) #对图像特征进行投影，得到图像嵌入向量
        image_embed = F.normalize(image_embed, dim=-1) #对图像嵌入向量进行归一化处理，使其在最后一个维度上具有单位长度

        vit_feats.append(vit_feat.cpu()) #将当前批次的vit特征添加到列表中，并移动到cpu上
        image_embeds.append(image_embed) #将当前批次的图像嵌入向量添加到列表中

    vit_feats = torch.cat(vit_feats, dim=0) #将所有批次的 ViT 特征在第0维上拼接成一个完整的张量。
    image_embeds = torch.cat(image_embeds, dim=0) #将所有批次的图像嵌入向量拼接成一个完整的张量。

    sims_matrix = [] #存储图像到文本的相似性矩阵
    for image_embed in image_embeds: #遍历每个图像嵌入向量
        sim_q2t = image_embed @ text_embeds.t() #计算当前图像嵌入向量与所有文本嵌入向量的点积相似性
        sim_i2t, _ = sim_q2t.max(0) #在每个文本的相似性分数中取最大值，得到图像到文本的相似性分数
        sims_matrix.append(sim_i2t) #将相似性分数添加到列表中
    sims_matrix = torch.stack(sims_matrix, dim=0) #将列表中的相似性分数张量在第0维上堆叠成一个完整的相似性矩阵

    score_matrix_i2t = torch.full(
        (len(data_loader.dataset.image), len(texts)), -100.0
    ).to(model.device) #初始化一个图像到文本的得分矩阵，填充值为-100.0，并移动设备

    num_tasks = dist_utils.get_world_size() #获取分布式环境中的进程数量
    rank = dist_utils.get_rank() #获取当前进程的排名，用于分布式训练中的进程标识
    step = sims_matrix.size(0) // num_tasks + 1 #计算每个进程需要处理的相似性矩阵的行数
    start = rank * step #计算当前进程处理的起始行索引
    end = min(sims_matrix.size(0), start + step) #计算当前进程处理的结束行索引，确保不超过矩阵的总行数

    for i, sims in enumerate(
        metric_logger.log_every(sims_matrix[start:end], 50, header)
    ): #使用metric_logger每50个迭代记录一次日志，遍历当前进程需要处理的相似性矩阵部分
        topk_sim, topk_idx = sims.topk(k=k_test, dim=0)
        image_inputs = vit_feats[start + i].repeat(k_test, 1, 1).to(model.device) #获取当前图像对应的vit特征，并重复k_test次
        score = model.compute_itm(
            image_inputs=image_inputs,
            text_ids=text_ids[topk_idx],
            text_atts=text_atts[topk_idx],
        ).float() #调用模型的compute_itm方法计算图像-文本匹配的得分，并将其转换为浮点类型
        score_matrix_i2t[start + i, topk_idx] = score + topk_sim #将得分与topk相似性分数相加，并存储到得分矩阵中

    sims_matrix = sims_matrix.t() #将相似性矩阵转置，以便后续计算文本到图像的相似性
    score_matrix_t2i = torch.full(
        (len(texts), len(data_loader.dataset.image)), -100.0
    ).to(model.device) #初始化一个文本到图像的得分矩阵，填充值为-100.0，并移动设备

    step = sims_matrix.size(0) // num_tasks + 1  #重新计算每个进程需要处理的转置后的相似性矩阵的行数
    start = rank * step #计算当前进程处理的起始行索引
    end = min(sims_matrix.size(0), start + step) #计算当前进程处理的结束行索引

    for i, sims in enumerate(
        metric_logger.log_every(sims_matrix[start:end], 50, header)
    ): #使用metric_logger每50个迭代记录一次日志，遍历当前进程需要处理的转置后的相似性矩阵部分
        topk_sim, topk_idx = sims.topk(k=k_test, dim=0)
        image_inputs = vit_feats[topk_idx.cpu()].to(model.device) #获取tokp图像对应的vit特征
        score = model.compute_itm(
            image_inputs=image_inputs,
            text_ids=text_ids[start + i].repeat(k_test, 1),
            text_atts=text_atts[start + i].repeat(k_test, 1),
        ).float() #调用模型的compute_itm方法计算文本-图像匹配的得分，并将其转换为浮点类型
        score_matrix_t2i[start + i, topk_idx] = score + topk_sim #将得分与topk相似性分数相加，并存储到得分矩阵中

    if dist_utils.is_dist_avail_and_initialized(): #判断是否处于分布式环境中并且已经初始化
        dist.barrier() #调用分布式进程，确保所有进程都到达这一点，防止后续操作出现不同步的问题
        torch.distributed.all_reduce(
            score_matrix_i2t, op=torch.distributed.ReduceOp.SUM
        ) #对图像到文本的得分矩阵进行全局求和操作，将所有进程的结果汇总到每个进程上
        torch.distributed.all_reduce(
            score_matrix_t2i, op=torch.distributed.ReduceOp.SUM
        ) #对文本到图像的得分矩阵进行全局求和操作，将所有进程的结果汇总到每个进程上

    total_time = time.time() - start_time #计算评估所花费的总时间
    total_time_str = str(datetime.timedelta(seconds=int(total_time))) #将总时间转换为格式化的时间字符串
    logging.info("Evaluation time {}".format(total_time_str)) #记录一条信息日志，显示评估花费的时间

    return score_matrix_i2t.cpu().numpy(), score_matrix_t2i.cpu().numpy()
